import os
import sys
import cv2
import numpy as np

"""
功能概述：
- 将输入图片的物体边缘提取，并生成抽象化的形状版本：
  1) 轮廓线稿  2) 多边形近似  3) 凸包  4) 平滑抽象形状
- 可用于“抽象物体形状”的基础处理与表达

运行示例：
python -u openCVTest.py input.jpg output_dir
输出：在 output_dir 下生成四种抽象结果图
"""

def ensure_dir(p):
    os.makedirs(p, exist_ok=True)


def read_image(path: str):
    img = cv2.imread(path)
    if img is None:
        raise FileNotFoundError(f"无法读取图像: {path}")
    return img


def preprocess(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    return gray


def segment_edges(gray):
    # Canny 边缘
    edges = cv2.Canny(gray, 60, 120)
    # 形态学闭运算，连通边缘
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel, iterations=1)
    return edges


def extract_main_contour(edges):
    cnts, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not cnts:
        return None
    # 取最大面积轮廓
    main = max(cnts, key=cv2.contourArea)
    return main


def draw_contour_line(img_shape, contour):
    h, w = img_shape[:2]
    canvas = np.ones((h, w, 3), dtype=np.uint8) * 255
    cv2.drawContours(canvas, [contour], -1, (0, 0, 0), 2)
    return canvas


def polygon_approx(contour, epsilon_ratio=0.01):
    # 根据轮廓周长的比例设置近似精度
    peri = cv2.arcLength(contour, True)
    epsilon = epsilon_ratio * peri
    approx = cv2.approxPolyDP(contour, epsilon, True)
    return approx


def draw_polygon(img_shape, poly):
    h, w = img_shape[:2]
    canvas = np.ones((h, w, 3), dtype=np.uint8) * 255
    cv2.polylines(canvas, [poly], True, (0, 0, 0), 2)
    return canvas


def convex_hull(contour):
    hull = cv2.convexHull(contour)
    return hull


def draw_hull(img_shape, hull):
    h, w = img_shape[:2]
    canvas = np.ones((h, w, 3), dtype=np.uint8) * 255
    cv2.polylines(canvas, [hull], True, (0, 0, 0), 2)
    return canvas


def smooth_abstract_shape(img_shape, contour, smooth_iter=3):
    # 将轮廓填充为mask
    h, w = img_shape[:2]
    mask = np.zeros((h, w), dtype=np.uint8)
    cv2.drawContours(mask, [contour], -1, 255, thickness=cv2.FILLED)
    # 高斯平滑再阈值，得到更“圆润”的抽象形状
    for _ in range(smooth_iter):
        mask = cv2.GaussianBlur(mask, (7, 7), 0)
        _, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
    canvas = np.ones((h, w, 3), dtype=np.uint8) * 255
    canvas[mask > 0] = (0, 0, 0)
    return canvas


def save_image(out_dir, name, img):
    ensure_dir(out_dir)
    path = os.path.join(out_dir, name)
    cv2.imwrite(path, img)
    return path


def abstract_shapes(input_path: str, out_dir: str):
    img = read_image(input_path)
    gray = preprocess(img)
    edges = segment_edges(gray)
    contour = extract_main_contour(edges)
    if contour is None:
        raise RuntimeError("未检测到显著轮廓，尝试换图或调整参数")

    # 1) 轮廓线稿
    line_img = draw_contour_line(img.shape, contour)
    save_image(out_dir, "abstract_line.png", line_img)

    # 2) 多边形近似（可调 epsilon_ratio 控制抽象程度）
    poly = polygon_approx(contour, epsilon_ratio=0.02)
    poly_img = draw_polygon(img.shape, poly)
    save_image(out_dir, "abstract_polygon.png", poly_img)

    # 3) 凸包抽象
    hull = convex_hull(contour)
    hull_img = draw_hull(img.shape, hull)
    save_image(out_dir, "abstract_hull.png", hull_img)

    # 4) 平滑抽象形状（更具艺术化）
    smooth_img = smooth_abstract_shape(img.shape, contour, smooth_iter=3)
    save_image(out_dir, "abstract_smooth.png", smooth_img)

    return {
        "line": os.path.join(out_dir, "abstract_line.png"),
        "polygon": os.path.join(out_dir, "abstract_polygon.png"),
        "hull": os.path.join(out_dir, "abstract_hull.png"),
        "smooth": os.path.join(out_dir, "abstract_smooth.png"),
    }


# 新增：矩形候选 + GrabCut 的楼房表面提取
import math
from mpl_toolkits.mplot3d.art3d import Poly3DCollection


def _angle_between(v1, v2):
    v1 = v1 / (np.linalg.norm(v1) + 1e-6)
    v2 = v2 / (np.linalg.norm(v2) + 1e-6)
    cosang = np.clip(np.dot(v1, v2), -1.0, 1.0)
    ang = math.degrees(math.acos(cosang))
    return ang


def is_rectangle_like(poly, angle_tol=15):
    # 近似四边形且角度接近 90°
    if len(poly) != 4:
        return False
    pts = poly.reshape(-1, 2)
    vecs = []
    for i in range(4):
        v = pts[(i + 1) % 4] - pts[i]
        vecs.append(v.astype(np.float32))
    angles = []
    for i in range(4):
        ang = _angle_between(vecs[i], vecs[(i + 1) % 4])
        angles.append(ang)
    return all(abs(a - 90) <= angle_tol for a in angles)


def find_rectangle_candidate(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(gray, 80, 160)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel, iterations=2)

    cnts, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not cnts:
        return None, None

    h, w = img.shape[:2]
    best = None
    best_score = -1
    for c in cnts:
        area = cv2.contourArea(c)
        if area < (h * w) * 0.02:
            continue
        peri = cv2.arcLength(c, True)
        approx = cv2.approxPolyDP(c, 0.02 * peri, True)
        if is_rectangle_like(approx, angle_tol=20):
            # 评分：面积大 + 长边接近垂直/水平
            rect = cv2.boundingRect(approx)
            x, y, rw, rh = rect
            # 宽高比在合理范围（楼房立面更偏矩形）
            ratio = max(rw, rh) / (min(rw, rh) + 1e-6)
            ratio_penalty = abs(ratio - 1.5)  # 经验：非正方更可能是立面
            score = area - ratio_penalty * 500
            if score > best_score:
                best_score = score
                best = (approx, rect)

    if best is None:
        # 回退：用最大轮廓的外接矩形
        c = max(cnts, key=cv2.contourArea)
        peri = cv2.arcLength(c, True)
        approx = cv2.approxPolyDP(c, 0.02 * peri, True)
        rect = cv2.boundingRect(c)
        return approx, rect
    return best


def grabcut_with_rect(img, rect, iter_count=6):
    x, y, w, h = rect
    mask = np.zeros(img.shape[:2], np.uint8)
    bgdModel = np.zeros((1, 65), np.float64)
    fgdModel = np.zeros((1, 65), np.float64)
    cv2.grabCut(img, mask, (x, y, w, h), bgdModel, fgdModel, iterCount=iter_count, mode=cv2.GC_INIT_WITH_RECT)
    # 生成前景掩码
    mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8')
    return mask2


def refine_mask(mask):
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
    return mask


def extract_building_surface(input_path: str, out_dir: str):
    img = cv2.imread(input_path)
    if img is None:
        raise FileNotFoundError(f"无法读取图像: {input_path}")
    approx, rect = find_rectangle_candidate(img)
    if rect is None:
        raise RuntimeError("未找到合适的矩形候选，图像背景复杂或边界不明显")

    mask = grabcut_with_rect(img, rect, iter_count=6)
    mask = refine_mask(mask)

    # 可视化
    h, w = img.shape[:2]
    canvas = np.ones((h, w), dtype=np.uint8) * 255
    canvas[mask > 0] = 0
    ensure_dir(out_dir)
    cv2.imwrite(os.path.join(out_dir, "building_mask.png"), canvas)

    # 轮廓与多边形
    cnts, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if cnts:
        c = max(cnts, key=cv2.contourArea)
        hull = cv2.convexHull(c)
        poly_img = np.ones((h, w, 3), dtype=np.uint8) * 255
        cv2.polylines(poly_img, [hull], True, (0, 0, 0), 2)
        cv2.imwrite(os.path.join(out_dir, "building_polygon.png"), poly_img)

    # 统计像素面积与占比
    area_px = int(mask.sum())
    total_px = h * w
    coverage = area_px / (total_px + 1e-6)
    with open(os.path.join(out_dir, "building_stats.txt"), "w", encoding="utf-8") as f:
        f.write(f"像素面积: {area_px}\n")
        f.write(f"图像总像素: {total_px}\n")
        f.write(f"覆盖比例: {coverage:.4f}\n")
    return {
        "mask": os.path.join(out_dir, "building_mask.png"),
        "polygon": os.path.join(out_dir, "building_polygon.png"),
        "stats": os.path.join(out_dir, "building_stats.txt"),
    }


# 柱体抽象：拟合圆并以楼体高度作为柱体高度进行3D渲染
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


def extract_building_mask(img):
    approx, rect = find_rectangle_candidate(img)
    if rect is None:
        raise RuntimeError("未找到合适的矩形候选，图像背景复杂或边界不明显")
    mask = grabcut_with_rect(img, rect, iter_count=6)
    mask = refine_mask(mask)
    return mask


def fit_cylinder_params(mask):
    cnts, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not cnts:
        raise RuntimeError("未在掩码中检测到楼体连通域")
    c = max(cnts, key=cv2.contourArea)
    (cx, cy), r = cv2.minEnclosingCircle(c)
    x, y, w, h = cv2.boundingRect(c)
    return {
        "center": (float(cx), float(cy)),
        "radius_px": float(r),
        "height_px": float(h),
        "bbox": (int(x), int(y), int(w), int(h)),
    }


def render_cylinder(radius_px: float, height_px: float, out_path: str):
    # 使用像素单位直接渲染柱体，展示抽象体积感
    theta = np.linspace(0, 2 * np.pi, 80)
    z = np.linspace(0, height_px, 80)
    Theta, Z = np.meshgrid(theta, z)
    X = radius_px * np.cos(Theta)
    Y = radius_px * np.sin(Theta)

    fig = plt.figure(figsize=(5, 6))
    ax = fig.add_subplot(111, projection='3d')
    ax.plot_surface(X, Y, Z, rstride=1, cstride=1, color='#3A78C1', alpha=0.95, linewidth=0)
    ax.plot(radius_px * np.cos(theta), radius_px * np.sin(theta), 0, color='#1f1f1f')
    ax.plot(radius_px * np.cos(theta), radius_px * np.sin(theta), height_px, color='#1f1f1f')
    ax.set_box_aspect((1, 1, height_px / (radius_px + 1e-6)))
    ax.view_init(elev=20, azim=35)
    ax.axis('off')
    plt.tight_layout()
    fig.savefig(out_path, dpi=140)
    plt.close(fig)


def abstract_building_as_cylinder(input_path: str, out_dir: str):
    img = cv2.imread(input_path)
    if img is None:
        raise FileNotFoundError(f"无法读取图像: {input_path}")
    mask = extract_building_mask(img)
    params = fit_cylinder_params(mask)

    ensure_dir(out_dir)
    cyl_path = os.path.join(out_dir, "cylinder_render.png")
    render_cylinder(params["radius_px"], params["height_px"], cyl_path)

    with open(os.path.join(out_dir, "cylinder_stats.txt"), "w", encoding="utf-8") as f:
        f.write(f"中心: {params['center']}\n")
        f.write(f"半径像素: {params['radius_px']:.2f}\n")
        f.write(f"高度像素: {params['height_px']:.2f}\n")
        x, y, w, h = params['bbox']
        f.write(f"外接矩形: (x={x}, y={y}, w={w}, h={h})\n")
    return {"cylinder": cyl_path, "stats": os.path.join(out_dir, "cylinder_stats.txt")}


# 建筑群候选检测（基于边缘 + 形态学 + 轮廓筛选）

# 改进版：垂直边缘驱动的候选检测 + 宽连通域自动切割

def _vertical_edge_mask(gray):
    gx = cv2.Sobel(gray, cv2.CV_16S, 1, 0, ksize=3)
    gx = cv2.convertScaleAbs(gx)
    # 归一化并二值化（Otsu）
    gx = cv2.normalize(gx, None, 0, 255, cv2.NORM_MINMAX)
    _, binx = cv2.threshold(gx, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    # 强化竖直结构
    k_vert = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 17))
    k_horiz = cv2.getStructuringElement(cv2.MORPH_RECT, (7, 1))
    m = cv2.morphologyEx(binx, cv2.MORPH_CLOSE, k_vert, iterations=2)
    m = cv2.morphologyEx(m, cv2.MORPH_OPEN, k_horiz, iterations=1)  # 抑制横向粘连
    return m


def _split_wide_rect_by_gaps(mask, rect, min_gap_px=6, min_width_px=20):
    x, y, w, h = rect
    roi = mask[y:y+h, x:x+w]
    col = roi.sum(axis=0)
    # 归一化列密度
    maxv = float(col.max()) + 1e-6
    norm = col / maxv
    # 找到低密度列（空隙），作为切割点
    gap_idxs = np.where(norm < 0.05)[0]
    if len(gap_idxs) == 0:
        return [rect]
    # 合并连续 gap 为区间
    splits = []
    start = gap_idxs[0]
    prev = gap_idxs[0]
    for idx in gap_idxs[1:]:
        if idx == prev + 1:
            prev = idx
        else:
            splits.append((start, prev))
            start = idx
            prev = idx
    splits.append((start, prev))
    # 选择足够宽的 gap
    cut_positions = [x + int((a + b) / 2) for (a, b) in splits if (b - a + 1) >= min_gap_px]
    if not cut_positions:
        return [rect]
    # 根据切割点生成子矩形
    xs = [x] + cut_positions + [x + w]
    subrects = []
    for i in range(len(xs) - 1):
        sx, ex = xs[i], xs[i + 1]
        sw = ex - sx
        if sw >= min_width_px:
            subrects.append((sx, y, sw, h))
    return subrects if subrects else [rect]


def find_building_candidates_v2(img, min_area_ratio=0.002, min_h_ratio=0.06, overlap_iou=0.45):
    h, w = img.shape[:2]
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    vm = _vertical_edge_mask(gray)
    # 连通域分析
    num, labels = cv2.connectedComponents(vm)
    rects = []
    img_area = h * w
    for lab in range(1, num):
        ys, xs = np.where(labels == lab)
        if ys.size == 0:
            continue
        y0, y1 = int(ys.min()), int(ys.max())
        x0, x1 = int(xs.min()), int(xs.max())
        rw, rh = x1 - x0 + 1, y1 - y0 + 1
        area = rw * rh
        if area < img_area * min_area_ratio:
            continue
        if rh < int(h * min_h_ratio):
            continue
        rects.append((x0, y0, rw, rh))
    # 对过宽的连通域进行自动切割
    split_rects = []
    for r in rects:
        x, y, rw, rh = r
        if rw > int(0.25 * w):  # 经验：超过图像宽度 1/4 的候选可能包含多栋
            subrs = _split_wide_rect_by_gaps(vm, r, min_gap_px=6, min_width_px=18)
            split_rects.extend(subrs)
        else:
            split_rects.append(r)

    # NMS 去重
    split_rects.sort(key=lambda r: r[2] * r[3], reverse=True)
    kept = []
    def iou(a, b):
        ax, ay, aw, ah = a
        bx, by, bw, bh = b
        x1 = max(ax, bx)
        y1 = max(ay, by)
        x2 = min(ax + aw, bx + bw)
        y2 = min(ay + ah, by + bh)
        inter = max(0, x2 - x1) * max(0, y2 - y1)
        ua = aw * ah + bw * bh - inter
        return inter / (ua + 1e-6)
    for r in split_rects:
        if all(iou(r, k) < overlap_iou for k in kept):
            kept.append(r)
    return kept

# 替换 cluster 模式使用 v2 检测

def abstract_building_cluster_as_cuboids(input_path: str, out_dir: str, depth_ratio=0.6):
    img = cv2.imread(input_path)
    if img is None:
        raise FileNotFoundError(f"无法读取图像: {input_path}")
    h, w = img.shape[:2]
    rects = find_building_candidates_v2(img)
    if not rects:
        raise RuntimeError("未检测到建筑群候选，可尝试放宽面积/高度阈值或更换图片")

    cuboids = []
    for (x, y, rw, rh) in rects:
        cuboids.append({
            'x': float(x),
            'y': float((y / h) * 0.5 * w),
            'w': float(rw),
            'd': float(max(10, int(rw * depth_ratio))),
            'h': float(rh),
        })

    ensure_dir(out_dir)
    out_render = os.path.join(out_dir, "cluster_cuboids.png")
    render_cuboids(cuboids, out_render)

    vis = img.copy()
    for (x, y, rw, rh) in rects:
        cv2.rectangle(vis, (x, y), (x + rw, y + rh), (0, 255, 0), 2)
    cv2.imwrite(os.path.join(out_dir, "cluster_rects.png"), vis)

    with open(os.path.join(out_dir, "cluster_stats.txt"), "w", encoding="utf-8") as f:
        f.write(f"建筑候选数量: {len(rects)}\n")
        for i, (x, y, rw, rh) in enumerate(rects):
            f.write(f"[{i}] x={x}, y={y}, w={rw}, h={rh}\n")
    return {"render": out_render, "rects_vis": os.path.join(out_dir, "cluster_rects.png"), "stats": os.path.join(out_dir, "cluster_stats.txt")}


# 更新 main：支持 "cluster" 模式
def main():
    if len(sys.argv) < 3:
        print("用法: python -u openCVTest.py <输入图片路径> <输出目录> [building|cylinder|cluster]")
        sys.exit(1)
    input_path = sys.argv[1]
    out_dir = sys.argv[2]
    mode = sys.argv[3].lower() if len(sys.argv) >= 4 else ""
    if mode == "building":
        res = extract_building_surface(input_path, out_dir)
        print("已生成楼房表面数据：")
        for k, v in res.items():
            print(f"- {k}: {v}")
    elif mode == "cylinder":
        res = abstract_building_as_cylinder(input_path, out_dir)
        print("已生成柱体抽象与统计：")
        for k, v in res.items():
            print(f"- {k}: {v}")
    elif mode == "cluster":
        res = abstract_building_cluster_as_cuboids(input_path, out_dir)
        print("已生成方柱体群抽象：")
        for k, v in res.items():
            print(f"- {k}: {v}")
    else:
        results = abstract_shapes(input_path, out_dir)
        print("已生成抽象形状：")
        for k, v in results.items():
            print(f"- {k}: {v}")


# 3D 方柱体渲染（提前定义，避免 main 调用前未定义）
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

def add_cuboid(ax, x, y, w, d, h, color="#4C9AFE", alpha=0.95):
    p0 = [x, y, 0];        p1 = [x + w, y, 0]
    p2 = [x + w, y + d, 0]; p3 = [x, y + d, 0]
    p4 = [x, y, h];        p5 = [x + w, y, h]
    p6 = [x + w, y + d, h]; p7 = [x, y + d, h]
    faces = [
        [p0, p1, p2, p3],
        [p4, p5, p6, p7],
        [p0, p1, p5, p4],
        [p1, p2, p6, p5],
        [p2, p3, p7, p6],
        [p3, p0, p4, p7],
    ]
    poly3d = Poly3DCollection(faces, facecolors=color, linewidths=0.2, edgecolors="#1f1f1f", alpha=alpha)
    ax.add_collection3d(poly3d)


def render_cuboids(cuboids, out_path):
    if not cuboids:
        raise RuntimeError("未检测到建筑候选，无法渲染方柱体群")
    fig = plt.figure(figsize=(6, 6))
    ax = fig.add_subplot(111, projection='3d')
    cuboids = sorted(cuboids, key=lambda c: (c['x'], c['y']))
    for c in cuboids:
        add_cuboid(ax, c['x'], c['y'], c['w'], c['d'], c['h'], color="#4C9AFE")
    max_x = max(c['x'] + c['w'] for c in cuboids)
    max_y = max(c['y'] + c['d'] for c in cuboids)
    max_h = max(c['h'] for c in cuboids)
    ax.set_box_aspect((max_x, max_y, max_h))
    ax.view_init(elev=20, azim=35)
    ax.axis('off')
    plt.tight_layout()
    fig.savefig(out_path, dpi=140)
    plt.close(fig)

if __name__ == "__main__":
    main()

def add_cuboid(ax, x, y, w, d, h, color="#4C9AFE", alpha=0.95):
    p0 = [x, y, 0];        p1 = [x + w, y, 0]
    p2 = [x + w, y + d, 0]; p3 = [x, y + d, 0]
    p4 = [x, y, h];        p5 = [x + w, y, h]
    p6 = [x + w, y + d, h]; p7 = [x, y + d, h]
    faces = [
        [p0, p1, p2, p3],
        [p4, p5, p6, p7],
        [p0, p1, p5, p4],
        [p1, p2, p6, p5],
        [p2, p3, p7, p6],
        [p3, p0, p4, p7],
    ]
    poly3d = Poly3DCollection(faces, facecolors=color, linewidths=0.2, edgecolors="#1f1f1f", alpha=alpha)
    ax.add_collection3d(poly3d)


def render_cuboids(cuboids, out_path):
    if not cuboids:
        raise RuntimeError("未检测到建筑候选，无法渲染方柱体群")
    fig = plt.figure(figsize=(6, 6))
    ax = fig.add_subplot(111, projection='3d')
    cuboids = sorted(cuboids, key=lambda c: (c['x'], c['y']))
    for c in cuboids:
        add_cuboid(ax, c['x'], c['y'], c['w'], c['d'], c['h'], color="#4C9AFE")
    max_x = max(c['x'] + c['w'] for c in cuboids)
    max_y = max(c['y'] + c['d'] for c in cuboids)
    max_h = max(c['h'] for c in cuboids)
    ax.set_box_aspect((max_x, max_y, max_h))
    ax.view_init(elev=20, azim=35)
    ax.axis('off')
    plt.tight_layout()
    fig.savefig(out_path, dpi=140)
    plt.close(fig)